딥이데아 Deep Idea

[논문리뷰] Emerging Properties in Self-Supervised Vision Transformers

2021년 05월 14일 Tags:
딥러닝, 논문리뷰


프레임워크

framework DINO: Self-distillation with no labels.

방법

  • 같은 아키텍쳐를 가진 Student model $g_{\theta_s}$과 Teacher model $g_{\theta_t}$을 준비한다.
    Student model은 gradient descent를 통해 업데이트하고,
    Teacher model은 Student model의 exponential moving average(ema)를 통해 업데이트한다.
  • 주어진 이미지 $\mathrm{x}$에 대하여 두 augmentation view의 쌍 $(\mathrm{x}_1, \mathrm{x}_2)$를 만들고,
    각각 다른 모델에 입력해준다.
    예를 들면, $\mathrm{x}_1$은 Student model에 $\mathrm{x}_2$는 Teachre model에 입력해준다.
  • 두 모델은 마지막으로 softmax layer를 통과한다.
    Teacher model의 output은 gradient가 흐르지 않게(sg; stop-gradient)
    고정된 텐서로 만들고, Student model의 output과 cross-entropy를 만들어 gradient descent해준다.
    예를 들면, $\mathrm{x}_1$이 입력된 Student model의 output을 $\mathrm{p}_1$이라 하고,
    $\mathrm{x}_2$가 입력된 Teacher model의 output을 $\mathrm{p}_2$라고 한다면,
    loss는 $-\mathrm{p}_2 \log \mathrm{p}_1$이다.
  • 즉, DINO는 Teacher model의 output을 마치 ground truth처럼 생각하고,
    Student model의 output과 cross-entropy를 구해 학습한다.
    Teacher model에서 Student model로 정보를 추출(distill)하는 방식의
    Self-supervised learning(SSL)이다.
    프레임워크의 이름 DINO는 Self-distillation with no labels에서 따온 것이다.
  • 안정적인 학습을 위해 centering과 sharpening이라는 테크닉을 추가했고,
    다른 SSL 방식에 비해 직관적이고,
    batch size가 커야 학습 효과가 좋은 contrastive learning 방식보다 효율적이다.

결과

  • Feature extractor를 담당하는 아키텍쳐는 ResNet50, DeiT, ViT를 이용했고,
    여러 SSL 방법과 ImageNet에서 비교 실험했다.
  • 평가 방법은 Linear evaluation(feature extractor 부분을 고정하고 output layer만 새로 학습하여 평가하는 방식)과 $k$-NN(representation을 가지고 $k$-NN으로 분류하는 방식)을 이용했다.

result Linear and k-NN classification on ImageNet

  • 위 표에서 보듯 DINO는 뛰어난 representation learning 능력을 보여주었다.
    특히 $k$-NN은 fine-tuning 없이 측정했음에도 뛰어난 성능을 보여주었다.
  • 또한 $k$-NN을 통해 좋은 퍼포먼스를 내기 위한 3가지 요소를 설명했는데,
    MoCo 논문에서 제시한 momentum encoder,
    SwAV 논문에서 제시한 multi-crop,
    그리고 본 논문에서 ViT를 사용할 때 작은 크기의 patch의 중요성을 제시했다.
    DINO는 이 3가지 요소를 모두 반영하여 학습했다.

map Self-attention map from Self-supervised learning with DINO

  • 이 논문에서 또 하나 강조한 것는 ViT를 SSL 했을 때
    마지막 block의 self-attention map이 위 그림처럼
    semantic segmentation에 대한 explicit한 정보를 담고 있다는 것이다.
    이것이 의미가 큰 이유는 supervision으로 학습한 ViT의 attention map에서도 explicit하게 나타나지 않던 특징이, SSL을 통해 배웠다는 것이다.
  • Self-attention의 각 head를 하나의 색깔(클래스)로 나타내면
    위 그림과 같이 실제 segmentation 처럼 visualize 할 수 있다.

Pseudocode

code DINO PyTorch pseudocode

  • loss 부분을 보면 $\mathrm{x}_1$과 $\mathrm{x}_2$를 Teacher model과 Student model에 교차입력해서 두 augmentation 순서쌍에 대해서 모두 학습해주는 것을 알 수 있다.
  • Student model을 SGD로, Teacher model을 ema로 업데이트 한다.
    ema의 momentum rate $l$은 cosine scheduling을 이용한다.
  • 학습 과정에서 collapse를 피하기 위해 Centering과 Sharpening을 사용하는데, 위 코드에는 C와 tps/tpt로 나타나있다.
    Centering은 output의 한 dimension이 dominate하는 것을 막지만,
    모든 output이 uniform하게 만드는 성질이 있고,
    Sharpening은 정반대의 성질이 있어, 두 가지를 함께 사용하여 균형을 맞춘다.
  • 이 방법은 batch에 대한 의존성을 낮추어 collapse를 피하는 것이다.
    Centering은 first-order batch statistics에 의존하고,
    Teacher model에 bias term을 더하는 것처럼 $g_t(x) \leftarrow g_t(x) + c$ 해석될 수 있다.
  • $c$는 다음과 같은 수식의 ema로 업데이트 되며 배치사이즈에 영향을 받을 수 있다.
    $ c \leftarrow mc + (1-m)\frac{1}{B} \sum_{i=1}^B g_{\theta_s}(x_i) $
    $m$은 hyperparameter다.
  • Sharpening은 작은 값의 temperature parameter tps/tpt를 사용하여
    softmax output을 normalize 해준다.
    Implementation을 살펴보면 Student model의 tps는 0.1,
    Teacher model의 tpt는 0.04에서 0.07로 linear warm-up을 해주는
    hyperparameter다.
  • Hyperparameter가 은근 많은 것 같다.

Reference

  1. Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, Armand Joulin
    Emerging Properties in Self-Supervised Vision Transformers
    https://arxiv.org/abs/2104.14294

Recent Posts